#导入包

from tensorflow.keras import Model
from tensorflow.keras.layers import Flatten, Conv2D, MaxPool2D, Dense, BatchNormalization, Dropout

"""
 tf.keras.layers.Conv2D
 参数：
 -input_shape ： 数据维数大小 (64,64,3)
                只在第一个网络层设置
 -filters : 输入卷积核个数，（输出后的通道数）
 -kernel_size: 卷积核的大小，通常为一个数或者一个列表
 例子：kernel_size = 3 kernel_size(3,) ，都表示卷积核的宽高都是为3
 -strides: 卷积核的滑动步长，通常是一个整数或者一个列表
 例子：strides=2,strides(2,2),都表示宽高的方向的滑动步长为2
 -padding : 补零设置 ，可选项"VALID","SAME"
 如果 padding="VALID",表示不做补零操作
 如果 padding="SAME" 程序自动补零，使得输出feature map 的宽高
 ceil(输入图片宽高/strides),ceil表示向上取整，当strides=1时，输入
 和输出宽高是一致的
 -activation: 激活函数
 可选项：
 -sigmoid:
 -tanh:
 -relu
 -elu 或者 tf.keras.activations.els(alpha=1.0)
 -selu 
 -swish : 在tf2.2版本以上才有
 -softmax
 -input_shape ：如是第一层卷积，需要设置输入图片的大小
 （height,width,channel),
 比如input_shape(128,128,3)
 -name 输入的字符串，代表该层设置一个名称
 
 池化操作 tf.keras.layers.Maxpool2D
 参数：
 -pool_size: 池化卷积核的大小，具体设置与Conv2D的kernel_size一样
 -strides :池化的滑动步长
 -name
 
 全连接操作 tf.keras.layers.Dense
 参数:
 -units : 全连接神经元个数
 -activation : 激活函数，输出层一般都用 "softmax"
 -name
 
 展平操作 tf.keras.layers.Flatten
 举例：
 [[1,2,3],[4,5,6],[7,8,9]]
 -->[1,2,3,4,5,6,7,8,9]
 
"""

# 定义一个子类


class ConvModel(Model):
    def __init__(self):
        super(ConvModel, self).__init__()

        self.conv_1_1 = Conv2D(
            # input_shape=(160, 160, 3),
            input_shape=(64, 64, 3),
            filters=32,
            kernel_size=3,
            activation='relu',
            name='conv_1_1'
        )

        self.bn_1_1 = BatchNormalization(axis=-1)
        self.dropout_1_1 = Dropout(0.5)

        self.conv_1_2 = Conv2D(
            filters=32,
            kernel_size=3,
            padding="SAME",
            activation='relu',
            name='conv_1_2'
        )

        self.bn_1_2 = BatchNormalization(axis=-1)
        self.dropout_1_2 = Dropout(0.5)

        self.conv_1_3 = Conv2D(
            filters=64,
            kernel_size=3,
            padding="SAME",
            activation='relu',
            name='conv_1_3'
        )

        self.bn_1_3 = BatchNormalization(axis=-1)
        self.dropout_1_3 = Dropout(0.5)

        self.max_pool_1 = MaxPool2D(
            pool_size=2,
            name='max_pool_1'
        )

        self.conv_2_1 = Conv2D(
            filters=64,
            kernel_size=3,
            padding="SAME",
            activation='relu',
            name='conv_2_1'
        )

        self.bn_2_1 = BatchNormalization(axis=-1)
        self.dropout_2_1 = Dropout(0.5)

        self.conv_2_2 = Conv2D(
            filters=64,
            kernel_size=3,
            padding="SAME",
            activation='relu',
            name='conv_2_2'
        )

        self.bn_2_2 = BatchNormalization(axis=-1)
        self.dropout_2_2 = Dropout(0.5)

        self.conv_2_3 = Conv2D(
            filters=64,
            kernel_size=3,
            padding="SAME",
            activation='relu',
            name='conv_2_3'
        )

        self.bn_2_3 = BatchNormalization(axis=-1)
        self.dropout_2_3 = Dropout(0.5)

        # 池化层
        self.max_pool_2 = MaxPool2D(
            pool_size=2,
            name='max_pool_2'
        )


        # 展平层
        self.flatten = Flatten(name='flatten')


        # 全连接层

        self.dense = Dense(
            units=10,
            activation='softmax',
            name='logit'
        )

    def call(self, x):
        x = self.conv_1_1(x)
        x = self.bn_1_1(x)
        #x = self.dropout_1_1(x)

        x = self.conv_1_2(x)
        x = self.bn_1_2(x)
        #x = self.dropout_1_2(x)

        x = self.conv_1_3(x)

        x = self.bn_1_3(x)
        #x = self.dropout_1_3(x)

        x = self.max_pool_1(x)

        x = self.conv_2_1(x)
        x = self.bn_2_1(x)
        #x = self.dropout_2_1(x)

        x = self.conv_2_2(x)
        x = self.bn_2_2(x)
        #x = self.dropout_2_2(x)

        x = self.conv_2_3(x)

        x = self.bn_2_3(x)

        x = self.dropout_2_3(x)

        x = self.max_pool_2(x)

        x = self.flatten(x)

        x = self.dense(x)

        return x

#model = ConvModel()
